Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@bionemo-recipes/models/qwen/collator.py`:
- Around line 975-977: The current calculation of slice_sizes uses floor
division (slice_sizes = (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) //
total_slices_of_any_sequence) which silently drops tokens when sequence lengths
are not divisible by total_slices_of_any_sequence (2 * cp_world_size); update
the logic around total_slices_of_any_sequence / slice_sizes / cu_seqlens_padded
to handle remainders instead of truncating — compute the base size and remainder
for each sequence (e.g., base = length // total_slices_of_any_sequence; rem =
length % total_slices_of_any_sequence) and distribute the rem extra tokens
across the first rem slices (or use ceil semantics and adjust offsets) so that
sum(slice_sizes) == original sequence length for every entry; ensure any
downstream code that consumes slice_sizes (referenced by
total_slices_of_any_sequence, slice_sizes, cu_seqlens_padded) uses these
adjusted slice sizes.
- Around line 435-437: The code sets max_length using the full unsharded
sequence length (max_length = batch["input_ids"].shape[1]) in the elif
self.qkv_format == "bshd" branch, which overstates per-CP-shard
max_length_q/max_length_k; change this to use the shard-local sequence length
instead: compute a shard_seq_len (either from an already-sharded tensor in
batch, e.g. a per-shard slice like batch["input_ids_shard"] if available, or by
dividing batch["input_ids"].shape[1] by self.cp_shards) and assign max_length =
shard_seq_len (and use that for max_length_q/max_length_k calculations) so each
CP shard uses its local sequence length in collator.py where qkv_format ==
"bshd".
- Around line 1015-1017: The function _scatter_batch_to_cp_tp_ranks currently
forbids None for all_batches but callsites sometimes pass None (e.g., where
combined_batch = ... if self.cp_tp_rank == 0 else None), so update the type
annotations to accept None and reflect that None can be returned: change the
parameter type to allow None (e.g., list[BatchType] | list[StopIteration] | None
or Optional[list[...]] | Optional[list[StopIteration]]) and adjust the return
type to BatchType | StopIteration | None; update any related type hints/imports
so Pyright no longer flags the callsite passing None.

In `@bionemo-recipes/models/qwen/convert_qwen2.py`:
- Around line 244-248: The TE→HF conversion currently always calls
output_model.tie_weights() and unconditionally uses model_hf._tied_weights_keys
for state_dict_ignored_entries, which breaks roundtrips when tie_word_embeddings
is False; update the TE→HF code around output_model.tie_weights() and the
state_dict_ignored_entries assignment to check
model_hf.config.tie_word_embeddings and only call output_model.tie_weights() and
set state_dict_ignored_entries = model_hf._tied_weights_keys when that flag is
True, otherwise leave state_dict_ignored_entries as an empty list and skip tying
weights so the behavior matches the HF→TE branch.
- Around line 75-87: The current code builds index tensors on CPU and then calls
.cpu() on the resulting bias slices (q_bias/k_bias/v_bias), which forces CUDA
qkv_bias onto CPU; instead, construct q_slice, k_slice, v_slice on the same
device as qkv_bias (e.g., torch.arange(..., device=qkv_bias.device)) and remove
the trailing .cpu() calls so q_bias, k_bias, v_bias remain on qkv_bias.device;
update the lines that create q_slice/k_slice/v_slice and the assignments to
q_bias/k_bias/v_bias in convert_qwen2.py accordingly (symbols: q_slice, k_slice,
v_slice, qkv_bias, q_bias, k_bias, v_bias, heads_per_group, num_query_groups,
qkv_total_dim).

In `@bionemo-recipes/models/qwen/export.py`:
- Line 56: The copy uses a relative path, so make it robust by resolving the
source file relative to this module: compute src = Path(__file__).parent /
"modeling_qwen3_te.py" (import Path from pathlib if needed) and pass that src to
shutil.copy(src, export_path / "modeling_qwen3_te.py"); update the copy call in
export.py to use this resolved Path instead of the relative string.

In `@bionemo-recipes/models/qwen/modeling_qwen2_te.py`:
- Around line 208-216: The code assumes input_ids exists when computing
padded_seq_len and when handling packing/unpacking; guard those accesses in the
should_pack_inputs path (and the other similar block around the
_unpad_input/_pad_input calls). Specifically, in the forward path where
should_pack_inputs is True and where you later reference input_ids.size(1),
replace direct input_ids usage with a conditional: if input_ids is not None use
input_ids.size(1) otherwise derive the sequence length from
inputs_embeds/hidden_states (e.g., hidden_states.size(1)); apply the same guard
for the other block around the _pad_input/_unpad_input calls so
input-only-embeds inputs do not raise AttributeError. Ensure you adjust variable
name padded_seq_len and any logic that uses it (e.g., the call sites of
_unpad_input/_pad_input) to use the guarded value.

In `@bionemo-recipes/models/qwen/modeling_qwen3_te.py`:
- Around line 141-143: The code uses config.layer_types[layer_idx] without
guarding for missing layer_types or a None sliding_window; update the
conditional that sets window_size (the expression around
window_size=(config.sliding_window, config.sliding_window) if
config.layer_types[layer_idx] == "sliding_attention" else None) to first verify
that config.layer_types exists and has an entry at layer_idx (e.g.,
hasattr/config.layer_types is not None and len(config.layer_types) > layer_idx)
and also check config.sliding_window is not None before treating the layer as
"sliding_attention"; if any check fails, set window_size to None.

In `@bionemo-recipes/models/qwen/state.py`:
- Around line 167-181: The code currently prints and continues when a target
parameter name is missing from target_state, which can mask unpopulated weights;
in the loop over target.named_parameters() (using variables name, param and the
target_state dict), replace the print(...) branch with a hard failure: raise a
descriptive ValueError (or RuntimeError) that includes the missing parameter
name and possibly remaining keys in target_state so conversion coverage is
clear, ensuring that _params is only populated when nn.Parameter(target_param,
requires_grad=param.requires_grad) is created and preventing silent
continuation.
- Around line 437-442: When there are no wildcard_matches the code uses
np.empty(shape, dtype=object) which leaves slots uninitialized; replace that
with a deterministic initialization (e.g. use np.full(shape, [], dtype=object)
or create an array populated with an explicit empty value) so the
`_match_keys`/output_array is reproducibly an empty container for no-match
cases; update the block that checks len(wildcard_matches) == 0 (variables:
wildcard_matches, shape, output_array, np.empty) to allocate and fill the array
deterministically.

In `@bionemo-recipes/models/qwen/tests/common/fixtures.py`:
- Around line 40-52: The fixture use_te_debug currently unconditionally pops
NVTE_DEBUG on teardown and thus can erase pre-existing environment values;
update use_te_debug (and any other TE-related fixtures that set env vars) to
save the original value (orig = os.environ.get("NVTE_DEBUG")) before setting it,
then on teardown restore it by setting it back if orig is not None or deleting
it only if orig was None; apply the same pattern for other fixtures in this file
that modify env vars so tests leave the environment exactly as they found it.

In `@bionemo-recipes/models/qwen/tests/test_modeling_qwen2_te.py`:
- Around line 68-73: The tokenizer is being loaded without the pinned revision,
causing a mismatch with the pinned model; update get_tokenizer to call
AutoTokenizer.from_pretrained with the same pinned revision by passing
revision=self.get_upstream_model_revision() (keep the existing pad_token
fallback logic), referencing the get_tokenizer method and
get_upstream_model_revision helper so the tokenizer and model use identical
revisions for reproducible tests.

In `@ci/scripts/check_copied_files.py`:
- Line 41: Update the destination paths in the list inside check_copied_files.py
that currently point to "bionemo-recipes/models/qwen3/..." (example entry
"bionemo-recipes/models/qwen3/collator.py") to the correct
"bionemo-recipes/models/qwen/..." paths (also fix the other two occurrences at
the other entries). This will ensure the validation list of expected
destinations (the list containing "bionemo-recipes/models/qwen3/collator.py")
matches the actual added files and prevents the ValueError raised by the
copied-files check.

---

Nitpick comments:
In `@bionemo-recipes/models/qwen/Dockerfile`:
- Around line 1-5: Add creation of a non-root user and switch to it in the
Dockerfile: create a user/group (e.g., bionemo) and set HOME, chown the WORKDIR
(/workspace/bionemo) and any cache directories (e.g., /root/.cache or change pip
cache path) to that user before switching; update RUN steps like the pip install
to run as the non-root user and add a final USER bionemo line. Ensure the USER
creation and chown occur after base image setup but before running application
commands so functions referencing WORKDIR, RUN
--mount=type=cache,target=/root/.cache/pip, and COPY . . will run under the
non-root account.
- Line 3: The Dockerfile uses a broad COPY . . which pulls the entire context;
update the repository .dockerignore to exclude build artifacts and VCS files so
they are not copied into the image. Add entries such as .git/, __pycache__/,
*.pyc, .venv/, venv/, build/, dist/, .pytest_cache/, .DS_Store and node_modules/
(as applicable) to the .dockerignore to shrink the build context and prevent
unintended files from being included.

In `@bionemo-recipes/models/qwen/modeling_qwen3_te.py`:
- Line 220: Replace the runtime-unsafe assertion with an explicit exception:
instead of using assert attention_mask is not None in the function/method that
validates BSHD inputs (the code referencing attention_mask), raise a ValueError
(or appropriate exception) with the same message ("Attention mask is required
when packing BSHD inputs.") so the check cannot be skipped under Python -O and
will always enforce the required input.
- Line 380: The module currently mutates global Dynamo state by setting
torch._dynamo.config.capture_scalar_outputs = True at import time; instead,
remove the top-level assignment and scope the change only where needed (e.g.,
around the code path that invokes compilation/tracing). Replace the global set
with a scoped approach such as using
torch._dynamo.config.patch(capture_scalar_outputs=True) (or save the original
value and restore it around the critical function) in the function that performs
tracing/compilation in modeling_qwen3_te.py so other modules are not affected by
a global config change.

In `@bionemo-recipes/models/qwen/requirements.txt`:
- Around line 1-5: The requirements.txt for the qwen recipe lists unpinned
dependencies (lm-eval, torch, torchao, transformer_engine[pytorch],
transformers) which breaks reproducibility; fix by either adding a
pyproject.toml for this recipe model that pins or bounds versions consistent
with other pinned projects (e.g., match amplify/geneformer ranges) or add
explicit version constraints/upper/lower bounds for each package in
requirements.txt (include torch, torchao, transformer_engine[pytorch],
transformers, lm-eval) to align with the project's reproducibility policy.

In `@bionemo-recipes/models/qwen/tests/conftest.py`:
- Around line 22-25: Remove the brittle sys.path.append manipulations in
conftest.py and instead import the local fixtures module directly to avoid
cross-package resolution: delete the two sys.path.append(...) calls, replace the
string-based pytest_plugins = ["tests.common.fixtures"] with an explicit local
import (e.g. from .common import fixtures) and set pytest_plugins = [fixtures]
(or import tests.common.fixtures as fixtures if relative import is unsuitable);
this keeps the unique module object (fixtures) referenced and prevents
accidental resolution of other packages named tests.common.fixtures.

In `@bionemo-recipes/models/qwen/tests/test_modeling_qwen3_te.py`:
- Around line 64-66: Replace the short commit hash in
get_upstream_model_revision with the full 40-character commit SHA to ensure
deterministic model loading; locate the get_upstream_model_revision() method and
change the returned string "c1899de" to the repository's full commit SHA for
that revision.
- Around line 38-41: The test file uses bare module imports that can break when
PYTHONPATH isn't set; replace them with explicit relative imports to the package
root so the test runs reliably. Update the import statements for
DataCollatorWithFlattening, convert_qwen3_hf_to_te, convert_qwen3_te_to_hf,
HFInferenceParams, NVQwen3Config, NVQwen3ForCausalLM, BaseModelTest, and
TestTolerances to use relative import syntax (e.g., from .collator import
DataCollatorWithFlattening or from ..models.convert_qwen3 import
convert_qwen3_hf_to_te) matching the test module's package level.
- Around line 134-143: Add a short explanatory comment above the get_tolerances
method describing why the chosen tolerances (golden_value_logits_atol=2.0 and
cp_loss_atol=0.5, cp_loss_rtol=0.25) are permissive — e.g., note they were
derived empirically from Qwen3 numerical variability or account for expected
floating-point/noise characteristics — so future readers understand the
rationale; update the docstring or add an inline comment referencing
get_tolerances and the TestTolerances fields to make the source of these
specific numeric choices clear.